import random
import itertools
import heapq

class PriorityQueue(object):
    def __init__(self):
        self.counter = itertools.count(1)
        self.queue = []
        self.index = {}
        
        heapq.heapify(self.queue)
    
    def pop(self):
        while True:
            priority, count, item = heapq.heappop(self.queue)
            del self.index[item]
            if count != 0:
                return item
    
    def add(self, item, priority):
        count = next(self.counter)
        
        entry = [priority, count, item]
        
        self.index[item] = entry
        
        heapq.heappush(self.queue, entry)
    
    def change(self, item, priority):
        if item in self.index:
            entry = self.index[item]
            entry[1] = 0
            
        self.add(item, priority)
    
    def __bool__(self):
        if not self.queue:
            return False
        
        while self.queue[0][1] == 0:
            del self.index[self.queue[0][2]]
            heapq.heappop(self.queue)
        
        if self.queue:
            return True
        
        return False

class PathingAI(object):
    """Maintain an index of good paths for units to follow."""
    
    def __init__(self, graph, objective):
        self.graph = graph
        self.objective = objective
        
        self.weights = {
            "test": {node: random.random() for node in graph.nodes}}
        self.paths = {}
    
    def calculate_paths(self):
        for label, weights in self.weights.iteritems():
            self.paths[label] = self.dijkstras(self.objective, weights)
    
    def dijkstras(self, sink, weights):
        paths = {}
        distances = {}
        
        for node in self.graph.nodes:
            distances[node] = 99999 #change this to something like infinity
            paths[node] = None
        
        paths[sink] = sink
        
        distances[sink] = 0
        settled = set()
        
        queue = PriorityQueue()
        queue.add(sink, 0)
        
        while queue.__bool__():
            u = queue.pop()
            
            settled.add(u)

            for v in u.neighbours:
                if distances[v] > distances[u] + weights[v]:
                    distances[v] = distances[u] + weights[v]
                    queue.change(v, distances[v])
                    paths[v] = u
        
        return paths
